import numpy as np
import pygame
from Box2D import b2_pi
import torch


def manhattan(x, y):
    return np.sum(np.abs(x - y))


def changeTheWorldByAction(action, worldContext):
    bodiesInTheWorld = worldContext['bodiesInTheWorld']
    dynamic_body = bodiesInTheWorld['dynamic_body']
    if action == 0:
        dynamic_body.position -= (0.1, 0)
    elif action == 1:
        dynamic_body.position += (0.1, 0)
    elif action == 2:
        dynamic_body.position += (0, 0.1)
    elif action == 3:
        dynamic_body.position -= (0, 0.1)
    elif action == 4:
        dynamic_body.angle += 0.1 * b2_pi
    elif action == 5:
        dynamic_body.angle -= 0.1 * b2_pi


def encodeRGB(rgb):
    result = rgb[0] + rgb[1] * 255 + rgb[2] * 255 * 255
    return result


def encodeImageTorch(img):
    imgT = torch.from_numpy(np.array(img)).type(torch.float).cuda(0)
    rT = imgT[:, :, 0]
    gT = imgT[:, :, 1]
    bT = imgT[:, :, 2]
    result = rT * 255 * 255 + gT * 255 + bT
    return result


def encodeCoord(coord, width):
    h, w, _ = coord
    order = h * width + w
    return order


def encodeImage(img):
    h, w, _ = img.shape
    result = np.zeros((h, w, 1), dtype='int32')
    for i in range(h):
        for j in range(w):
            result[i, j] = encodeRGB(img[i, j])
    return result
